package cool.mtc.security.plugin.wrapper;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.util.*;

/**
 * @author yz
 */
public class RequestHeaderWrapper extends HttpServletRequestWrapper {
    private final Map<String, String> header = new HashMap<>();

    public RequestHeaderWrapper(HttpServletRequest request) {
        super(request);
    }

    public void addHeader(String key, String value) {
        header.put(key, value);
    }

    @Override
    public String getHeader(String name) {
        if (header.containsKey(name)) {
            return header.get(name);
        }
        return super.getHeader(name);
    }

    @Override
    public Enumeration<String> getHeaders(String name) {
        if (header.containsKey(name)) {
            List<String> list = Collections.list(super.getHeaders(name));
            list.add(header.get(name));
            return Collections.enumeration(list);
        }
        return super.getHeaders(name);
    }

    @Override
    public Enumeration<String> getHeaderNames() {
        if (!header.isEmpty()) {
            List<String> list = Collections.list(super.getHeaderNames());
            list.addAll(header.keySet());
            return Collections.enumeration(list);
        }
        return super.getHeaderNames();
    }
}
